package com.jd.jobkeeper.utils;

import com.google.protobuf.CodedInputStream;
import com.google.protobuf.ExtensionRegistry;
import com.google.protobuf.Message;
import io.rsocket.Payload;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class ProtobufDecoder {

    private final ExtensionRegistry extensionRegistry;

    private static final ConcurrentMap<Class<?>, Method> methodCache = new ConcurrentHashMap<>();

    /**
     * Construct a new {@code ProtobufDecoder}.
     */
    public ProtobufDecoder() {
        this(ExtensionRegistry.newInstance());
    }

    /**
     * Construct a new {@code ProtobufDecoder} with an initializer that allows the
     * registration of message extensions.
     * @param extensionRegistry a message extension registry
     */
    public ProtobufDecoder(ExtensionRegistry extensionRegistry) {
        Objects.requireNonNull(extensionRegistry, "ExtensionRegistry must not be null");
        this.extensionRegistry = extensionRegistry;
    }

    public boolean canDecode(Class elementType) {
        return Message.class.isAssignableFrom(elementType);
    }



    public Message decode(Payload payload,Class<?> clazz)  {

        try {
            Message.Builder builder = getMessageBuilder(clazz);
            builder.mergeFrom(CodedInputStream.newInstance(payload.getData()), this.extensionRegistry);
            return builder.build();
        }
        catch (IOException ex) {
            ex.printStackTrace();
            //throw new DecodingException("I/O error while parsing input stream", ex);
        }
        catch (Exception ex) {
            //throw new DecodingException("Could not read Protobuf message: " + ex.getMessage(), ex);
        }
        finally {
            payload.release();
        }
        return  null;
    }

    /**
     * Create a new {@code Message.Builder} instance for the given class.
     * <p>This method uses a ConcurrentHashMap for caching method lookups.
     */
    private static Message.Builder getMessageBuilder(Class<?> clazz) throws Exception {
        Method method = methodCache.get(clazz);
        if (method == null) {
            method = clazz.getMethod("newBuilder");
            methodCache.put(clazz, method);
        }
        return (Message.Builder) method.invoke(clazz);
    }

}
